from models.meta import BaseDocument, BaseEmbeddedDocument, fields


class NameType(BaseEmbeddedDocument):
    zh_cn = fields.StringField()
    en_us = fields.StringField()


class KbVuln(BaseDocument):
    meta = {'db_alias': 'kb', 'collection': 'kb_vuln', 'strict': False}

    source = fields.StringField(required=True)
    source_id = fields.StringField(required=True)
    name = fields.EmbeddedDocumentField(document_type=NameType, required=True)
    vuln_type = fields.EmbeddedDocumentListField(document_type=NameType)
    is_plugin = fields.BooleanField(default=True)
    updated_at = fields.DateTimeField()
    created_at = fields.DateTimeField()
    published_at = fields.DateTimeField()
    severity = fields.StringField(default='')
    description = fields.EmbeddedDocumentField(document_type=NameType)
    cnnvd = fields.ListField(default=[])
    cve = fields.ListField(default=[])
    impact = fields.EmbeddedDocumentField(document_type=NameType)
    owasp = fields.ListField(default=[])
    recommendation = fields.EmbeddedDocumentField(document_type=NameType)
    reference = fields.ListField(default=[])
    zoomeye_dork = fields.StringField(default='')
    zoomeye_count = fields.IntField(default=0)
    is_deleted = fields.BooleanField(default=False)


class SystemKeywords(BaseDocument):
    meta = {'db_alias': 'kb', 'collection': 'system_keywords', 'strict': False}

    word = fields.StringField(required=True, max_length=100)
    type = fields.IntField(required=True, default=8, min_value=1, max_value=8)
    level = fields.IntField(required=False, default=3, min_value=1, max_value=3)
    is_disabled = fields.BooleanField(required=False, default=False)
    is_deleted = fields.BooleanField(required=False, default=False)
    template_id = fields.StringField(required=False, default='default')


class CustomKeywords(BaseDocument):
    meta = {'db_alias': 'kb', 'collection': 'custom_keywords', 'strict': False}

    word = fields.StringField(required=True, max_length=100)
    type = fields.IntField(required=True, default=8, min_value=1, max_value=8)
    level = fields.IntField(required=False, default=3, min_value=1, max_value=3)
    is_disabled = fields.BooleanField(required=False, default=False)
    is_deleted = fields.BooleanField(required=False, default=False)
    custom_id = fields.StringField(required=True)


class WhiteListKeywords(BaseDocument):
    meta = {'db_alias': 'kb', 'collection': 'whitelist_keywords', 'strict': False}

    word = fields.StringField(required=True, max_length=100)
    custom_id = fields.StringField(required=True)
